package com.leimo.web_socket2.interceptor;

import com.leimo.common.Constant;

import org.apache.commons.lang3.StringUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.util.Map;

import javax.servlet.http.HttpSession;

/**
 * WebSocket握手请求的拦截器. 检查握手请求和响应
 * @author 段誉
 * @create 2019-02-26 14:58
 */
public class CustomHandshakeInterceptor extends HttpSessionHandshakeInterceptor {
  /**
   * 在握手之前执行该方法, 继续握手返回true, 中断握手返回false. 通过attributes参数设置WebSocketSession的属性
   * 还有个afterHandshake方法，在握手之后执行该方法. 无论是否握手成功都指明了响应状态码和相应头.
   * @param request
   * @param response
   * @param wsHandler
   * @param attributes
   * @return
   * @throws Exception
   */
  @Override
  public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                 WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

    //获取sessionId
    ServletServerHttpRequest httpRequest = (ServletServerHttpRequest) request;
    String sessionId = httpRequest.getServletRequest().getParameter(Constant.WEB_SOCKET_SESSION_ID);
    HttpSession session = httpRequest.getServletRequest().getSession();

    if (StringUtils.isNoneEmpty(sessionId)) {
      session.setAttribute(Constant.WEB_SOCKET_SESSION_ID, sessionId);
    }

    return super.beforeHandshake(request, response, wsHandler, attributes);
  }
}
